#pragma once
#include <cinttypes>

#include "macro.h"

template <typename T_int>
HOST_DEVICE_INLINE T_int div_up(T_int a, T_int b) {
  return (a + b - 1) / b;
}

/// Object to encapsulate the fast division+modulus operation.
///
/// This object precomputes two values used to accelerate the computation and is
/// best used when the divisor is a grid-invariant. In this case, it may be
/// computed in host code and marshalled along other kernel arguments using the
/// 'Params' pattern.
///
/// Example:
///
///
///   int quotient, remainder, dividend, divisor;
///
///   FastDivmod divmod(divisor);
///
///   divmod(quotient, remainder, dividend);
///
///   // quotient = (dividend / divisor)
///   // remainder = (dividend % divisor)
///
struct FastDivmod {
  int divisor;
  unsigned int multiplier;
  unsigned int shift_right;

  template <typename value_t>
  HOST_DEVICE_INLINE static value_t clz(value_t x) {
    for (int i = 31; i >= 0; --i) {
      if ((1 << i) & x) return value_t(31 - i);
    }
    return 32;
  }

  template <typename value_t>
  HOST_DEVICE_INLINE static value_t find_log2(value_t x) {
    int a = int(31 - clz(x));
    a += (x & (x - 1)) != 0;  // Round up, add 1 if not a power of 2.
    return a;
  }

  /// Find quotient and remainder using device-side intrinsics
  HOST_DEVICE_INLINE
  void fast_divmod(int &quotient, int &remainder, int dividend) const {
#if defined(__CUDA_ARCH__)
    // Use IMUL.HI if divisor != 1, else simply copy the source.
    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right
                              : dividend;
#else
    quotient =
        int((divisor != 1)
                ? int(((int64_t)dividend * multiplier) >> 32) >> shift_right
                : dividend);
#endif

    // The remainder.
    remainder = dividend - (quotient * divisor);
  }

  /// For long int input
  HOST_DEVICE_INLINE
  void fast_divmod(int &quotient, int64_t &remainder, int64_t dividend) const {
#if defined(__CUDA_ARCH__)
    // Use IMUL.HI if divisor != 1, else simply copy the source.
    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right
                              : dividend;
#else
    quotient =
        int((divisor != 1) ? ((dividend * multiplier) >> 32) >> shift_right
                           : dividend);
#endif
    // The remainder.
    remainder = dividend - (quotient * divisor);
  }

  /// Construct the FastDivmod object, in host code ideally.
  ///
  /// This precomputes some values based on the divisor and is computationally
  /// expensive.

  HOST_DEVICE_INLINE
  FastDivmod() : divisor(0), multiplier(0), shift_right(0) {}

  HOST_DEVICE_INLINE
  FastDivmod(int divisor) : divisor(divisor) {
    if (divisor != 1) {
      unsigned int p = 31 + find_log2(divisor);
      unsigned m =
          unsigned(((1ull << p) + unsigned(divisor) - 1) / unsigned(divisor));

      multiplier = m;
      shift_right = p - 32;
    } else {
      multiplier = 0;
      shift_right = 0;
    }
  }

  /// Computes integer division and modulus using precomputed values. This is
  /// computationally inexpensive.
  HOST_DEVICE_INLINE
  void operator()(int &quotient, int &remainder, int dividend) const {
    fast_divmod(quotient, remainder, dividend);
  }

  /// Computes integer division using precomputed values. This is
  /// computationally inexpensive.
  HOST_DEVICE_INLINE
  int div(int dividend) const {
    int quotient, remainder;
    fast_divmod(quotient, remainder, dividend);
    return quotient;
  }

  /// Alias for `div` to match the interface of FastDivmodU64
  HOST_DEVICE_INLINE
  int divide(int dividend) const { return div(dividend); }

  /// Computes integer division and modulus using precomputed values. This is
  /// computationally inexpensive.
  ///
  /// Simply returns the quotient
  HOST_DEVICE_INLINE
  int divmod(int &remainder, int dividend) const {
    int quotient;
    fast_divmod(quotient, remainder, dividend);
    return quotient;
  }

  /// Computes integer division and modulus using precomputed values. This is
  /// computationally inexpensive.
  HOST_DEVICE_INLINE
  void operator()(int &quotient, int64_t &remainder, int64_t dividend) const {
    fast_divmod(quotient, remainder, dividend);
  }

  /// Computes integer division and modulus using precomputed values. This is
  /// computationally inexpensive.
  HOST_DEVICE_INLINE
  int divmod(int64_t &remainder, int64_t dividend) const {
    int quotient;
    fast_divmod(quotient, remainder, dividend);
    return quotient;
  }

  /// Returns the divisor when cast to integer
  HOST_DEVICE_INLINE
  operator int() const { return divisor; }
};
